#!/usr/bin/env python3
import numpy as np
from typing import Dict, Any
import sys
import os
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent / 'RoboTwin'))
sys.path.append(str(Path(__file__).parent.parent / 'debugger'))
from envs.utils.vis_pcd import keypose48_to_pcd, catpcd
from save_file import save_info
from scalePcd import *
from scipy.spatial.transform import Rotation

def get_48d_key_poses_definition():
    """
    定义48D key_poses的结构
    
    Returns:
        dict: 包含48D key_poses结构定义的字典
        
    48D key_poses = [12D contact_poses + 36D gripper_poses]
    
    12D contact_poses (前12维):
        - [0:3]: 第一个接触点的XYZ位置
        - [3:6]: 第一个接触点的RPY旋转 (roll, pitch, yaw)
        - [6:9]: 第二个接触点的XYZ位置  
        - [9:12]: 第二个接触点的RPY旋转 (roll, pitch, yaw)
    
    36D gripper_poses (后36维):
        - [12:15]: 左gripper的XYZ位置
        - [15:18]: 左gripper的RPY旋转 (roll, pitch, yaw)
        - [18:21]: 左gripper左手指attach点的XYZ位置
        - [21:24]: 左gripper左手指end点的XYZ位置
        - [24:27]: 左gripper右手指attach点的XYZ位置
        - [27:30]: 左gripper右手指end点的XYZ位置
        - [30:33]: 右gripper的XYZ位置
        - [33:36]: 右gripper的RPY旋转 (roll, pitch, yaw)
        - [36:39]: 右gripper左手指attach点的XYZ位置
        - [39:42]: 右gripper左手指end点的XYZ位置
        - [42:45]: 右gripper右手指attach点的XYZ位置
        - [45:48]: 右gripper右手指end点的XYZ位置
    """
    return {
        "contact_1": {"pos": [0, 1, 2], "rpy": [3, 4, 5]},
        "contact_2": {"pos": [6, 7, 8], "rpy": [9, 10, 11]},
        "left_gripper": {
            "pos": [12, 13, 14],
            "rpy": [15, 16, 17],
            "left_finger_attach": [18, 19, 20],
            "left_finger_end": [21, 22, 23],
            "right_finger_attach": [24, 25, 26],
            "right_finger_end": [27, 28, 29]
        },
        "right_gripper": {
            "pos": [30, 31, 32],
            "rpy": [33, 34, 35],
            "left_finger_attach": [36, 37, 38],
            "left_finger_end": [39, 40, 41],
            "right_finger_attach": [42, 43, 44],
            "right_finger_end": [45, 46, 47]
        }
    }

def keypose_item(keypose48D: np.ndarray) -> Dict[str, Any]:
    """
    解析48D key_poses数组，返回包含实际值的字典
    
    Args:
        keypose48D: 48维的key_poses数组
        
    Returns:
        dict: 包含实际值的字典，结构与get_48d_key_poses_definition()相同
        
    Example:
        keypose48D = np.random.rand(48)
        result = parse_48d_key_poses(keypose48D)
        print(result['contact_poses']['contact_1']['pos'])  # [x, y, z] values
        print(result['gripper_poses']['left_gripper']['pos'])  # [x, y, z] values
    """
    # if 
    assert keypose48D.shape == (48,), f"Expected keypose48D shape (48,), got {keypose48D.shape}"
    
    # Get the structure definition
    structure_def = get_48d_key_poses_definition()
    
    def extract_values_from_structure(structure_dict, data_array):
        """Recursively extract values from structure definition using data array"""
        result = {}
        for key, value in structure_dict.items():
            if isinstance(value, dict):
                result[key] = extract_values_from_structure(value, data_array)
            elif isinstance(value, list):
                # Extract values from data array using the indices in the list
                result[key] = [data_array[i] for i in value]
        return result
    
    return extract_values_from_structure(structure_def, keypose48D)

def extendContactPose(key_poses: np.ndarray, dist: float, idx: int) -> np.ndarray:
    """Extend contact pose by distance in RPY direction. idx=0: first XYZ, idx=6: second XYZ."""
    assert key_poses.shape == (48,), f"Expected key_poses shape (48,), got {key_poses.shape}"
    extended = key_poses.copy()
    xyz, rpy = extended[idx:idx+3], extended[idx+3:idx+6]
    rotation = Rotation.from_euler('xyz', rpy)
    direction = rotation.apply([0, 0, 1])  # Z-axis direction after rotation
    extended[idx:idx+3] = xyz + direction * dist
    return extended

def keyPcd(unifiedData: Dict[str, Any], dist: float = 0.0, *, surpress_print: bool = False) -> Dict[str, Any]:
    assert 'key_poses' in unifiedData and unifiedData['key_poses'] is not None, "No key_poses data available for keyPcd"
    key_poses = unifiedData['key_poses']
    assert key_poses.shape == (48,), f"Expected key_poses shape (48,), got {key_poses.shape}"
    
    # Extend contact poses if distance > 0
    if dist > 0:
        key_poses = extendContactPose(key_poses, dist, 0) # Extend first contact pose
        key_poses = extendContactPose(key_poses, dist, 6) # Extend second contact pose
        unifiedData['key_poses'] = key_poses  # Update in place
        if not surpress_print:
            print(f"INFO: Extended contact poses by {dist} distance")
    
    # Convert 48D key_poses to point cloud using keypose48_to_pcd
    #test: override the two rpy to 0

    key_pcd = keypose48_to_pcd(key_poses)
    assert key_pcd.shape[1] == 6 and key_pcd.shape[0] > 0, f"Expected key_pcd shape (N, 6), got {key_pcd.shape}"
    
    # Scale colors to 0-255 range using scalePcd
    # key_pcd = scalePcd(key_pcd, tag="key_pcd")
    
    if not surpress_print:
        print(f"INFO: Generated key point cloud with {key_pcd.shape[0]} points")
    
    updated_fields = {}
    for pcd_key in unifiedData.keys():
        if "pointcloud" in pcd_key and unifiedData[pcd_key] is not None:
            existing_pcd = unifiedData[pcd_key]
            assert len(existing_pcd.shape) == 2 and existing_pcd.shape[1] == 6, f"Skipping {pcd_key}: Expected shape (Np, 6), got {existing_pcd.shape}"
            original_count = existing_pcd.shape[0]
            # --- COLOR NORMALIZATION ---
            existing_pcd = scalePcd(existing_pcd, tag="existing_pcd")
            key_pcd_to_use = scalePcd(key_pcd, tag="key_pcd_to_use")
            combined_pcd = catpcd(existing_pcd, key_pcd_to_use)
            unifiedData[pcd_key] = combined_pcd
            updated_fields[pcd_key] = (original_count, combined_pcd.shape[0])
    
    for key, (original_count, final_count) in updated_fields.items():
        print(f"    🔑 KEYPCD({key}): {original_count} → {final_count}")
    
    return unifiedData

if __name__ == "__main__":
    test_data = {
        'pointcloud': np.random.rand(100, 6),
        'key_poses': np.random.rand(48),  # 48D key_poses instead of separate contact_poses and gripper_poses
        'rgb': None,
        'cameras': {},
        'gaussian': None
    }
    result = keyPcd(test_data, dist=0.1)
    print(f"Test completed. Point cloud shape: {result['pointcloud'].shape}")